/*
 * To change this license header, choose License Headers in Project Properties.
 * To change this template file, choose Tools | Templates
 * and open the template in the editor.
 */
package cn.ac.iie.di.commons.speeder.limiter;

import cn.ac.iie.di.commons.speeder.lifecycle.LifeCycleException;
import java.util.LinkedList;
import java.util.concurrent.locks.ReentrantLock;

/**
 *
 * @author zwlin
 */
public class LazyTokenBucket implements TokenBucket {

    private static class Pair<K, V> {

        public K key;
        public V value;

        public Pair(K key, V value) {
            this.key = key;
            this.value = value;
        }

    }

    private static final int THOUSAND = 1000;
    private static final int ZERO = 0;

    private boolean activated;
    private long max;
    private long deltaTimestamp;
    private long cur;
    private LinkedList<Pair<Long/*timestamp*/, Integer/*token num*/>> gots;

    ReentrantLock lock;

    public LazyTokenBucket(long tokenPerSec) {
        this(tokenPerSec, THOUSAND);
    }

    public LazyTokenBucket(long tokenEachTime, long intervalMsec) {
        max = tokenEachTime;
        deltaTimestamp = intervalMsec;
        init();
    }

    private void init() {
        activated = false;
        cur = ZERO;
        lock = new ReentrantLock();
        gots = new LinkedList<>();
    }

    @Override
    public boolean get(int num) {
        //not activated
        if (!activated) {
            throw new LifeCycleException("not started");
        }
        //unable to get
        if (max < num) {
            return false;
        }
        //wait until get
        lock.lock();
        try {
            //delete expired gots
            while (gots.size() > 0) {
                if (gots.getFirst().key > System.currentTimeMillis() - deltaTimestamp) {
                    break;
                }
                cur -= gots.removeFirst().value;
            }
            //return if get enough
            if (cur + num <= max) {
                cur += num;
                gots.add(new Pair<>(System.currentTimeMillis(), num));
                return true;
            }
            //wait until get enough
            do {
                if (!activated) {
                    return false;
                }
                try {
                    //calc wait time
                    long waitUntil = 0;
                    long need = cur + num - max;
                    for (Pair<Long, Integer> pair : gots) {
                        need -= pair.value;
                        if (need <= 0) {
                            waitUntil = pair.key;
                            break;
                        }
                    }
                    //calc
                    Thread.sleep(Math.max(ZERO, waitUntil + deltaTimestamp - System.currentTimeMillis()));
                    while (gots.size() > 0) {
                        if (gots.getFirst().key > System.currentTimeMillis() - deltaTimestamp) {
                            break;
                        }
                        cur -= gots.removeFirst().value;
                    }
                } catch (InterruptedException ex) {
                }

            } while (cur + num > max);
            cur += num;
            gots.add(new Pair<>(System.currentTimeMillis(), num));
            return true;
        } finally {
            lock.unlock();
        }
    }

    @Override
    public boolean tryGet(int num) {
        //not activated
        if (!activated) {
            throw new LifeCycleException("not started");
        }
        //unable to get
        if (max < num) {
            return false;
        }
        //calc if get
        lock.lock();
        try {
            //delete expired gots
            while (gots.size() > 0) {
                if (gots.getFirst().key > System.currentTimeMillis() - deltaTimestamp) {
                    break;
                }
                cur -= gots.removeFirst().value;
            }
            //return if get enough
            if (cur + num <= max) {
                cur += num;
                gots.add(new Pair<>(System.currentTimeMillis(), num));
                return true;
            } else {
                return false;
            }
        } finally {
            lock.unlock();
        }
    }

    @Override
    public void start() {
        activated = true;
    }

    @Override
    public void stop() {
        activated = false;
    }

}
